package GMDPRE

import "fmt"

// Sm4RotL 将src循环左移n位
func Sm4RotL(src uint32, n int) uint32 {
	n = n & 0x1f
	return src<<n | src>>(32-n)
}

// Sm4FL 轮函数使用的线性变换 L
func Sm4FL(src uint32) uint32 {
	return src ^ Sm4RotL(src, 2) ^ Sm4RotL(src, 10) ^ Sm4RotL(src, 18) ^ Sm4RotL(src, 24)
}

// Sm4Tau 非线性变换 Tau(τ)
func Sm4Tau(a uint32) uint32 {
	return (uint32(sBox[a>>24&0xFF]) << 24) |
		(uint32(sBox[a>>16&0xFF]) << 16) |
		(uint32(sBox[a>>8&0xFF]) << 8) |
		(uint32(sBox[a&0xFF]))
}

func bytes2Words(input []byte) []uint32 {
	var x = make([]uint32, Sm4WordSize)
	for i := 0; i < Sm4WordSize; i++ {
		x[i] = (uint32(input[4*i]) << 24) |
			(uint32(input[1+4*i]) << 16) |
			(uint32(input[2+4*i]) << 8) |
			uint32(input[3+4*i])
	}
	return x
}

func words2Bytes(words []uint32) []uint8 {
	var x = make([]uint8, Sm4BlockSize)
	for i := 0; i < Sm4WordSize; i++ {
		x[i*4] = uint8((words[i] >> 24) & 0xff)
		x[i*4+1] = uint8((words[i] >> 16) & 0xff)
		x[i*4+2] = uint8((words[i] >> 8) & 0xff)
		x[i*4+3] = uint8(words[i] & 0xff)
	}
	return x
}

func crypt(m []byte, rk []uint32) []uint8 {
	var X = bytes2Words(m)
	for i := 0; i < Sm4RkSize; i++ {
		// 轮函数
		X[i&0b11] = X[i&0b11] ^ Sm4FL(Sm4Tau(X[(i+1)&0b11]^X[(i+2)&0b11]^X[(i+3)&0b11]^rk[i]))
	}
	return words2Bytes([]uint32{X[3], X[2], X[1], X[0]})
}

// Sm4EL 密钥扩展时使用的线性变换 L'
func Sm4EL(src uint32) uint32 {
	return src ^ Sm4RotL(src, 13) ^ Sm4RotL(src, 23)
}

func keyExt(keyBytes []uint8) []uint32 {
	if len(keyBytes) != Sm4BlockSize {
		fmt.Printf("Key Error!!!!!")
		return []uint32{}
	}
	var MK = bytes2Words(keyBytes)
	var K = []uint32{MK[0] ^ FK[0], MK[1] ^ FK[1], MK[2] ^ FK[2], MK[3] ^ FK[3]}
	var rk = make([]uint32, Sm4RkSize)

	for i := 0; i < Sm4RkSize; i++ {
		K[i&0b11] = K[i&0b11] ^ Sm4EL(Sm4Tau(K[(i+1)&0b11]^K[(i+2)&0b11]^K[(i+3)&0b11]^CK[i]))
		rk[i] = K[i&0b11]
	}

	return rk
}

func Sm4Encrypt(mBytes []uint8, mLen int, keyBytes []uint8) []uint8 {
	var rk = keyExt(keyBytes)
	var t, k, outLen int
	t = mLen / Sm4BlockSize
	k = mLen % Sm4BlockSize
	outLen = t * Sm4BlockSize

	if k > 0 {
		outLen += Sm4BlockSize
	}

	var res = make([]uint8, outLen)
	for i := 0; i < t; i++ {
		copy(res[i*Sm4BlockSize:(i+1)*Sm4BlockSize], crypt(mBytes[i*Sm4BlockSize:(i+1)*Sm4BlockSize], rk))
	}

	if k > 0 {
		tmp := make([]uint8, Sm4BlockSize)
		copy(tmp[:k], mBytes[t*Sm4BlockSize:])
		copy(res[t*Sm4BlockSize:], crypt(tmp, rk))
	}

	return res
}

func Sm4Decrypt(cBytes []uint8, cLen int, keyBytes []uint8) []uint8 {
	var rk = keyExt(keyBytes)

	for i := 0; i < Sm4RkSize/2; i++ {
		j := Sm4RkSize - i - 1
		rk[i] ^= rk[j]
		rk[j] ^= rk[i]
		rk[i] ^= rk[j]
	}

	var t, k, outLen int
	t = cLen / Sm4BlockSize
	k = cLen % Sm4BlockSize
	outLen = cLen

	if k > 0 {
		fmt.Println("Sm4 Decrypt failed!!!!!!")
	}

	var res = make([]uint8, outLen)
	for i := 0; i < t; i++ {
		copy(res[i*Sm4BlockSize:(i+1)*Sm4BlockSize], crypt(cBytes[i*Sm4BlockSize:(i+1)*Sm4BlockSize], rk))
	}

	return res
}
